import os
import random

import PIL
import h5py
import imageio
import numpy
from PIL import Image
from icecream import ic

from data_pre.data import create_validate_code
from utils.coding import encode

data_img_path = 'data_img'
data_path = 'data_img.h5'
scale = 0.8


def generate_img(num):
    for i in range(num):
        img, label = create_validate_code()
        img.save(data_img_path + os.sep + label + '.gif')
        print(i / num * 100)


def read_img():
    img_list = os.listdir(data_img_path)
    random.shuffle(img_list)

    res = []
    labels = []

    i = 0
    total_len = len(img_list)
    for filename in img_list:
        i += 1
        print(i / total_len)

        try:
            img_array = PIL.Image.open(data_img_path + os.sep + filename)
            img_array = numpy.array(img_array)

            label = filename.split('.')[0]
            label = numpy.array(encode(label))

            res.append(img_array)
            labels.append(label)
        except Exception:
            ic(filename)

    return res, labels


def make_h5():
    print("making h5 file......")
    # Create a new file
    f = h5py.File(data_path, 'w')

    # 读取照片和对应的标签
    res, labels = read_img()
    data_len = int(len(res) * scale)

    f.create_dataset('x_train', data=res[:data_len])
    f.create_dataset('y_train', data=labels[:data_len])
    f.create_dataset('x_test', data=res[data_len:])
    f.create_dataset('y_test', data=labels[data_len:])
    f.close()


def read_data():
    print('reading data')

    datasets = h5py.File(data_path, 'r')
    x_train = numpy.array(datasets['x_train'][:])
    y_train = numpy.array(datasets['y_train'][:])
    x_test = numpy.array(datasets['x_test'][:])
    y_test = numpy.array(datasets['y_test'][:])
    datasets.close()

    x_train = numpy.expand_dims(x_train, axis=-1)
    x_test = numpy.expand_dims(x_test, axis=-1)

    ic(x_train.shape)
    ic(y_train.shape)
    ic(x_test.shape)
    ic(y_test.shape)

    return x_train, y_train, x_test, y_test
